import pytest
import tempfile
from load.custom import TaggedTextSplitter, DocImport
import asyncio

@pytest.fixture(name="short_text")
def short_text():
    # 模拟的测试文本
    return """
    稿件Start
    这是第一段内容。
    稿件End

    稿件Start
    这是第二段内容。
    稿件End
    """

@pytest.fixture(name="long_text")
def long_text():
    # 模拟的测试文本
    return """


稿件Start
稿件标题：提要2（1.1）
稿件日期：2024年01月01日
记者：1
摄像：1
配音：朱佳明
稿件内容：
 【正文】 
1、消防、文化执法、医护人员、地铁公交工作人员节日坚守岗位，全力守护市民安康。
【保安全护民生】2、元旦假期进入尾声，我市返程交通车流平稳。
【返程交通平稳】      
稿件End



稿件Start
稿件标题：我市召开第三十四次人大工作研讨会暨市人大工作研究会六届五次理事会议
稿件日期：2022年01月03日
记者：彭晔
摄像：陈华兴
配音：丁力
稿件内容：我市召开第三十四次人大工作研讨会暨市人大工作研究会六届五次理事会议
彭晔 陈华兴
 【口播】 
    来看一条要闻
 【正文】 
    近日（12.29），常州市第三十四次人大工作研讨会暨市人大工作研究会六届五次理事会议在市行政中心召开。
    会议强调，要毫不动摇坚持、与时俱进完善人民代表大会制度；要认真贯彻落实发展全过程人民民主的要求，确保党委政府在决策环节、执行环节、监督落实环节都能听到来自人民的声音；要加强理论研究和实践探索，以高质量的研究成果和切实可行的思路举措，为开创新时代地方人大工作新局面贡献智慧和力量。
    会议回顾总结了研究会一年来的工作和取得的成效，并对2022年研究会工作进行了部署。
    市人大常委会常务副主任张耀钢、副主任蒋自平和秘书长高宏华参加会议。
稿件End


    """

@pytest.fixture
def test_temp_dir():
    # 创建一个临时目录并返回路径
    with tempfile.TemporaryDirectory() as tmp_dir:
        yield tmp_dir

@pytest.mark.asyncio
async def test_tagged_text_splitter(short_text):
    # 测试 TaggedTextSplitter 的 split_text 方法
    splitter = TaggedTextSplitter()
    result = splitter.split_text(short_text)

    assert len(result) == 2  # 应该有两个分段
    assert result[0] == "这是第一段内容。"  # 第一个分段内容
    assert result[1] == "这是第二段内容。"  # 第二个分段内容

@pytest.mark.asyncio
async def test_doc_import_split_files(test_temp_dir, short_text):
    doc_import = DocImport(source_dir=test_temp_dir)

    documents = doc_import._load_single_file(
        "fakename.txt", 
        short_text, 
        TaggedTextSplitter())

    assert len(documents) == 2  # 应该有两个文档
    assert documents[0]["page_content"] == "这是第一段内容。"
    assert documents[1]["page_content"] == "这是第二段内容。"

def test_doc_import_extract_meta(test_temp_dir, long_text):
    doc_import = DocImport(source_dir=test_temp_dir)

    documents = documents = doc_import._load_single_file(
        "fakename.txt", 
        long_text, 
        TaggedTextSplitter())

    correct, incorrect = asyncio.run(doc_import._extract_meta(documents))
    assert len(correct) == 1
    assert len(incorrect) == 1

    assert correct[0]["metadata"]["title"] == "我市召开第三十四次人大工作研讨会暨市人大工作研究会六届五次理事会议"
    assert correct[0]["metadata"]["reporter"] == ["彭晔"]
    assert correct[0]["metadata"]["cameraman"] == ["陈华兴"]
    assert correct[0]["metadata"]["voiceover"] == ["丁力"]
    assert correct[0]["metadata"]["reported_on"] == "2022/01/03"
    